import torch
import time

# 基础信息输出
print("=== 环境信息 ===")
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA版本: {torch.version.cuda}")
print(f"CuDNN版本: {torch.backends.cudnn.version()}")
print(f"GPU可用: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"显卡型号: {torch.cuda.get_device_name(0)}")
    print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
else:
    print("未检测到可用的CUDA设备")

# 显存性能测试
print("\n=== 显存性能测试 ===")
if torch.cuda.is_available():
    torch.cuda.empty_cache()
    size = 512
    try:
        x = torch.randn(size, size, size, device='cuda')
        start = time.time()
        y = x + x
        torch.cuda.synchronize()  # 等待GPU完成
        print(f"矩阵运算耗时: {time.time() - start:.4f}秒")
        del x, y
    except Exception as e:
        print(f"测试过程中发生错误: {e}")
else:
    print("无法执行GPU性能测试: 没有可用的CUDA设备")